# coding:UTF-8
# 合并RTK航向角和IMU数据记录的程序
import binascii
import serial
import time
import csv
import re
import os
from datetime import datetime
import pynmea2
from pyproj import CRS, Transformer

# --- 串口参数定义 ---
# 电机控制串口 - 已注释
# motor_params = {
#     'port': 'COM6',  # 电机驱动板串口号
#     'baudrate': 115200,
#     'timeout': 0.5
# }

# GPS/RTK串口
rtk_params = {
    'port': '/dev/ttyUSB1',  # RTK GPS串口号
    'baudrate': 115200,
    'timeout': 0.5
}

# IMU串口
imu_params = {
    'port': '/dev/ttyUSB0',  # IMU串口号
    'baudrate': 9600,
    'timeout': 0.5
}

# --- 全局变量 ---
# IMU数据相关
ACCData = [0.0]*8
GYROData = [0.0]*8
AngleData = [0.0]*8
FrameState = 0  # 帧状态判断
Bytenum = 0  # 当前帧已接收字节数
CheckSum = 0  # 校验和

a = [0.0]*3  # 加速度
w = [0.0]*3  # 角速度
Angle = [0.0]*3  # 角度

# RTK/GPS数据相关
current_heading = None  # 当前航向角
rtk_timestamp = None  # RTK时间戳

# --- 电机控制相关函数 --- (保留但不使用)
def generate_control_command(speed_value, angle_deflection=0):
    """生成控制指令
    Args:
        speed_value: 速度值 (0-10000)
        angle_deflection: 角度值 (默认为0)
    Returns:
        str: 控制指令字符串
    """
    angle_hex = f"{angle_deflection & 0xFFFFFFFF:08X}"  # 角度，限制在32位无符号整数范围内
    speed_hex = f"{-speed_value & 0xFFFFFFFF:08X}"  # 速度，负值表示反向
    
    # 指令格式： "E0 01 00 00 [角度] [速度]"
    command = f"E0010000{angle_hex}{speed_hex}"
    return command

def send_command(ser, command):
    """发送控制指令到电机
    Args:
        ser: 串口对象
        command: 控制指令字符串
    """
    try:
        # 将十六进制字符串转换为字节
        byte_command = binascii.unhexlify(command)
        ser.write(byte_command)
        print(f"发送命令: {command}")
    except Exception as e:
        print(f"发送指令错误: {e}")

# --- RTK数据处理函数 ---
def process_hdt_data(data):
    """处理HDT语句，获取真实航向角
    Args:
        data: NMEA语句字符串
    Returns:
        float: 航向角
    """
    try:
        sentence = pynmea2.parse(data)
        if isinstance(sentence, pynmea2.HDT):
            true_heading = float(sentence.heading)
            return true_heading
    except (pynmea2.nmea.ChecksumError, pynmea2.ParseError):
        # 忽略ChecksumError和ParseError异常
        pass
    except Exception as e:
        # 捕获并打印其他异常
        print(f"处理HDT数据错误: {e}")
    return None

# --- IMU数据处理函数 ---
def DueData(inputdata):
    """处理IMU原始数据
    Args:
        inputdata: IMU发送的原始字节数据
    Returns:
        bool: 是否成功解析出一组完整数据
    """
    global FrameState
    global Bytenum
    global CheckSum
    global a
    global w
    global Angle
    
    for data in inputdata:  # 遍历输入数据
        if FrameState == 0:  # 状态0：寻找帧头
            if data == 0x55 and Bytenum == 0:  # 找到第一个帧头字节0x55
                CheckSum = data
                Bytenum = 1
                continue
            elif data == 0x51 and Bytenum == 1:  # 加速度数据帧
                CheckSum += data
                FrameState = 1
                Bytenum = 2
            elif data == 0x52 and Bytenum == 1:  # 角速度数据帧
                CheckSum += data
                FrameState = 2
                Bytenum = 2
            elif data == 0x53 and Bytenum == 1:  # 角度数据帧
                CheckSum += data
                FrameState = 3
                Bytenum = 2
        elif FrameState == 1:  # 状态1：接收加速度数据
            if Bytenum < 10:  # 读取8个数据字节
                ACCData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):  # 校验和检查
                    a = get_acc(ACCData)
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
        elif FrameState == 2:  # 状态2：接收角速度数据
            if Bytenum < 10:
                GYROData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):
                    w = get_gyro(GYROData)
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
        elif FrameState == 3:  # 状态3：接收角度数据
            if Bytenum < 10:
                AngleData[Bytenum-2] = data
                CheckSum += data
                Bytenum += 1
            else:
                if data == (CheckSum & 0xff):
                    Angle = get_angle(AngleData)
                    return True  # 返回True表示成功解析了一组完整数据
                CheckSum = 0
                Bytenum = 0
                FrameState = 0
    return False  # 返回False表示没有解析到完整数据

def get_acc(datahex):
    """计算加速度值
    Args:
        datahex: 加速度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的加速度
    """
    axl = datahex[0]
    axh = datahex[1]
    ayl = datahex[2]
    ayh = datahex[3]
    azl = datahex[4]
    azh = datahex[5]
    k_acc = 16.0
    acc_x = (axh << 8 | axl) / 32768.0 * k_acc
    acc_y = (ayh << 8 | ayl) / 32768.0 * k_acc
    acc_z = (azh << 8 | azl) / 32768.0 * k_acc
    if acc_x >= k_acc:
        acc_x -= 2 * k_acc
    if acc_y >= k_acc:
        acc_y -= 2 * k_acc
    if acc_z >= k_acc:
        acc_z -= 2 * k_acc
    return acc_x, acc_y, acc_z

def get_gyro(datahex):
    """计算角速度值
    Args:
        datahex: 角速度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的角速度
    """
    wxl = datahex[0]
    wxh = datahex[1]
    wyl = datahex[2]
    wyh = datahex[3]
    wzl = datahex[4]
    wzh = datahex[5]
    k_gyro = 2000.0
    gyro_x = (wxh << 8 | wxl) / 32768.0 * k_gyro
    gyro_y = (wyh << 8 | wyl) / 32768.0 * k_gyro
    gyro_z = (wzh << 8 | wzl) / 32768.0 * k_gyro
    if gyro_x >= k_gyro:
        gyro_x -= 2 * k_gyro
    if gyro_y >= k_gyro:
        gyro_y -= 2 * k_gyro
    if gyro_z >= k_gyro:
        gyro_z -= 2 * k_gyro
    return gyro_x, gyro_y, gyro_z

def get_angle(datahex):
    """计算角度值
    Args:
        datahex: 角度的原始数据字节
    Returns:
        tuple: (x, y, z)轴的角度
    """
    rxl = datahex[0]
    rxh = datahex[1]
    ryl = datahex[2]
    ryh = datahex[3]
    rzl = datahex[4]
    rzh = datahex[5]
    k_angle = 180.0
    angle_x = (rxh << 8 | rxl) / 32768.0 * k_angle
    angle_y = (ryh << 8 | ryl) / 32768.0 * k_angle
    angle_z = (rzh << 8 | rzl) / 32768.0 * k_angle
    if angle_x >= k_angle:
        angle_x -= 2 * k_angle
    if angle_y >= k_angle:
        angle_y -= 2 * k_angle
    if angle_z >= k_angle:
        angle_z -= 2 * k_angle
    return angle_x, angle_y, angle_z

def read_rtk_data(ser_rtk):
    """读取RTK数据，提取航向角
    Args:
        ser_rtk: RTK串口对象
    Returns:
        heading: 航向角
        timestamp: 时间戳
    """
    global current_heading, rtk_timestamp
    
    if ser_rtk.in_waiting > 0:
        data = ser_rtk.readline().decode('utf-8', errors='ignore').rstrip()
        heading = process_hdt_data(data)
        if heading is not None:
            current_heading = heading
            rtk_timestamp = datetime.now().strftime('%H:%M:%S.%f')
            return current_heading, rtk_timestamp
    
    return current_heading, rtk_timestamp

# --- 主程序 ---
if __name__ == '__main__':
    try:
        # 打开串口 - 电机串口已注释
        # ser_motor = serial.Serial(**motor_params)
        # print(f"电机串口已打开: {ser_motor.is_open}, 端口: {motor_params['port']}")
        
        ser_rtk = serial.Serial(**rtk_params)
        print(f"RTK串口已打开: {ser_rtk.is_open}, 端口: {rtk_params['port']}")
        
        ser_imu = serial.Serial(**imu_params)
        print(f"IMU串口已打开: {ser_imu.is_open}, 端口: {imu_params['port']}")
        
        # 获取用户输入 - 修改为不询问速度
        # speed_input = input("请输入电机速度值(0-10000)或'q'退出: ")
        # if speed_input.lower() == 'q':
        #     print("程序退出")
        #     exit(0)
        
        # 速度检查 - 已注释
        # if not re.match(r"^\d+$", speed_input):
        #     print("输入无效，应为0-10000的整数。默认设置为0。")
        #     speed = 0
        # else:
        #     speed = int(speed_input)
        #     if not 0 <= speed <= 10000:
        #         print("速度超出范围，应该在0-10000之间。默认设置为0。")
        #         speed = 0
        
        # 创建CSV文件 - 修改文件名
        filename_input = input("请输入保存文件的名称(如输入数字则保存为'rtk_imu_数字.csv'，输入'q'退出): ")
        if filename_input.lower() == 'q':
            print("程序退出")
            exit(0)
            
        filename = f"rtk_imu_{filename_input}.csv"
        print(f"将保存数据到文件: {filename}")
        
        with open(filename, 'w', newline='') as csvfile:
            csv_writer = csv.writer(csvfile)
            csv_writer.writerow([
                'Timestamp', 
                'RTK_Heading', 'RTK_Timestamp',
                'IMU_Acc_X', 'IMU_Acc_Y', 'IMU_Acc_Z', 
                'IMU_Gyro_X', 'IMU_Gyro_Y', 'IMU_Gyro_Z',
                'IMU_Angle_X', 'IMU_Angle_Y', 'IMU_Angle_Z'
                # 'Motor_Speed' - 已移除
            ])
            
            # 发送控制命令启动电机 - 已注释
            # print(f"启动电机，速度: {speed}")
            # motor_command = generate_control_command(speed)
            # send_command(ser_motor, motor_command)
            # time.sleep(1)  # 等待电机响应
            
            # 设置运行时间
            run_time = 5.0  # 修改为5秒
            start_time = time.time()
            data_count = 0
            
            print(f"开始记录数据，将运行{run_time}秒...")
            
            # 主循环
            while time.time() - start_time < run_time:
                # 读取IMU数据
                datahex = ser_imu.read(33)
                
                # 处理IMU数据
                if datahex and DueData(datahex):
                    # 显示IMU数据
                    print("加速度: %10.3f %10.3f %10.3f" % a)
                    print("角速度: %10.3f %10.3f %10.3f" % w)
                    print("角度: %10.3f %10.3f %10.3f" % Angle)
                    print("-" * 50)
                    
                    # 读取RTK数据
                    heading, rtk_time = read_rtk_data(ser_rtk)
                    
                    # 显示RTK数据
                    if heading is not None:
                        print(f"RTK航向角: {heading}°, 时间: {rtk_time}")
                    
                    # 记录数据 - 移除motor_speed
                    timestamp = time.time() - start_time
                    csv_writer.writerow([
                        timestamp,
                        heading, rtk_time,
                        a[0], a[1], a[2],
                        w[0], w[1], w[2],
                        Angle[0], Angle[1], Angle[2]
                        # speed - 已移除
                    ])
                    csvfile.flush()
                    data_count += 1
                    
                    if data_count % 10 == 0:
                        elapsed = time.time() - start_time
                        remaining = max(0, run_time - elapsed)
                        print(f"已记录 {data_count} 组数据, 剩余时间: {remaining:.1f}秒")
                
                # 继续读取RTK数据，即使没有IMU数据
                read_rtk_data(ser_rtk)
                
                # 短暂休眠以减少CPU使用率
                time.sleep(0.01)
            
            # 时间到，停止电机 - 已注释
            # print("运行时间结束，停止电机")
            # stop_command = "E00100000000000000000000"
            # send_command(ser_motor, stop_command)
            
            print(f"共记录 {data_count} 组数据")
            
            # 添加等待用户输入，让用户选择继续记录或退出
            while True:
                next_action = input("记录完成。输入新的文件名继续记录，或输入'q'退出: ")
                if next_action.lower() == 'q':
                    break
                    
                # 新文件名
                filename = f"rtk_imu_{next_action}.csv"
                print(f"将保存数据到新文件: {filename}")
                
                with open(filename, 'w', newline='') as csvfile:
                    csv_writer = csv.writer(csvfile)
                    csv_writer.writerow([
                        'Timestamp', 
                        'RTK_Heading', 'RTK_Timestamp',
                        'IMU_Acc_X', 'IMU_Acc_Y', 'IMU_Acc_Z', 
                        'IMU_Gyro_X', 'IMU_Gyro_Y', 'IMU_Gyro_Z',
                        'IMU_Angle_X', 'IMU_Angle_Y', 'IMU_Angle_Z'
                    ])
                    
                    # 重置计数和时间
                    start_time = time.time()
                    data_count = 0
                    
                    print(f"开始记录新数据，将运行{run_time}秒...")
                    
                    # 主循环
                    while time.time() - start_time < run_time:
                        # 读取IMU数据
                        datahex = ser_imu.read(33)
                        
                        # 处理IMU数据
                        if datahex and DueData(datahex):
                            # 显示IMU数据
                            print("加速度: %10.3f %10.3f %10.3f" % a)
                            print("角速度: %10.3f %10.3f %10.3f" % w)
                            print("角度: %10.3f %10.3f %10.3f" % Angle)
                            print("-" * 50)
                            
                            # 读取RTK数据
                            heading, rtk_time = read_rtk_data(ser_rtk)
                            
                            # 显示RTK数据
                            if heading is not None:
                                print(f"RTK航向角: {heading}°, 时间: {rtk_time}")
                            
                            # 记录数据
                            timestamp = time.time() - start_time
                            csv_writer.writerow([
                                timestamp,
                                heading, rtk_time,
                                a[0], a[1], a[2],
                                w[0], w[1], w[2],
                                Angle[0], Angle[1], Angle[2]
                            ])
                            csvfile.flush()
                            data_count += 1
                            
                            if data_count % 10 == 0:
                                elapsed = time.time() - start_time
                                remaining = max(0, run_time - elapsed)
                                print(f"已记录 {data_count} 组数据, 剩余时间: {remaining:.1f}秒")
                        
                        # 继续读取RTK数据，即使没有IMU数据
                        read_rtk_data(ser_rtk)
                        
                        # 短暂休眠以减少CPU使用率
                        time.sleep(0.01)
                    
                    print(f"共记录 {data_count} 组数据")
    
    except KeyboardInterrupt:
        print("\n程序被用户中断")
        # 注释以下电机停止代码
        # if 'ser_motor' in locals() and ser_motor.is_open:
        #     print("紧急停止电机")
        #     stop_command = "E00100000000000000000000"
        #     send_command(ser_motor, stop_command)
    
    except serial.SerialException as e:
        print(f"串口错误: {e}")
    
    except Exception as e:
        print(f"程序错误: {e}")
    
    finally:
        # 关闭所有串口 - 电机串口已注释
        for ser_name, ser_obj in [
            # ('ser_motor', 'ser_motor'), - 已注释
            ('ser_rtk', 'ser_rtk'), 
            ('ser_imu', 'ser_imu')
        ]:
            if ser_obj in locals() and locals()[ser_obj].is_open:
                locals()[ser_obj].close()
                print(f"{ser_name}串口已关闭") 